import os
from langchain_community.chat_models import ChatTongyi
from langchain_core.messages import HumanMessage, SystemMessage
from langchain.chains import create_sql_query_chain
from langchain_community.chat_models import ChatOpenAI  # 或者使用其他支持的LLM
from langchain_community.utilities import SQLDatabase
from langchain.tools import tool
from langchain.agents import AgentType, initialize_agent, create_tool_calling_agent
from langchain.agents import AgentExecutor



# 1、读取环境变量配置的api_key和url
AI_QW_KEY = os.getenv("AI_QW_KEY")
AI_QW_URL = os.getenv("AI_QW_URL")

# 设置通义千问的 API Key
os.environ["DASHSCOPE_API_KEY"] = AI_QW_KEY  # 替换为你的 API Key
# 可选：设置通义千问的 API 端点（如果需要）
os.environ["DASHSCOPE_ENDPOINT"] = AI_QW_URL

# 初始化通义千问模型
llm = ChatTongyi()

# 初始化数据库连接（假设库名为xeasy，根据你的实际情况修改连接参数）
db = SQLDatabase.from_uri(
    f"mysql+pymysql://root:root@127.0.0.1:3306/db_xeasy",  # 替换为你的连接信息
    sample_rows_in_table_info=2  # 在表信息中示例的行数
)

# 创建SQL查询链
chain = create_sql_query_chain(llm, db)

# 定义工具
@tool
def query_student_natural_language(question: str) -> str:
    """根据自然语言问题查询tf_task_step表。’"""
    try:
        # 让LLM生成SQL并执行
        sql_query = chain.invoke({"question": question})
        print(f"生成的SQL: {sql_query}")  # 调试用
        result = db.run(sql_query)
        return f"问题: '{question}' 的查询结果：{result}"
    except Exception as e:
        return f"查询失败：{str(e)}"



tools = [query_student_natural_language]


# 创建并执行Agent
# 使用 ZERO_SHOT_REACT_DESCRIPTION 代理类型
agent_executor = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,  # 使用零样本ReAct代理
    verbose=True,  # 显示详细执行过程，便于调试
    handle_parsing_errors=True  # 优雅处理解析错误
)
# 现在你可以使用 agent_executor 来运行你的查询
result = agent_executor.invoke({"input": "查询tf_task_step表有几条数据"})
print(result['output'])























# 定义消息列表
# messages = [
#     SystemMessage(content="你是一个地理老师。"),
#     HumanMessage(content="珠穆朗玛峰高度是多少")
# ]
#
# # 调用模型生成响应
# response = model.invoke(messages)
# print("模型响应:", response)
